package com.dayang.framework.biz.limiter.aspect;

import com.dayang.framework.biz.limiter.util.IpUtil;
import com.dayang.framework.common.exception.BizException;
import com.dayang.framework.common.exception.CommonResponseCodeEnum;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.util.Collections;
import java.util.List;

/**
 * @author: dayang
 * @since: 2024/6/13
 * @description: 限流切面类
 */
@Aspect
@Slf4j
public class RateLimiterAspect {
    private final RedisTemplate<String,Object> redisTemplate;
    private final DefaultRedisScript<Long> limitScript;

    public RateLimiterAspect(RedisTemplate<String, Object> redisTemplate, DefaultRedisScript<Long> limitScript) {
        this.redisTemplate = redisTemplate;
        this.limitScript = limitScript;
    }
    @Before("@annotation(rateLimiter)")
    public void before(JoinPoint joinPoint,RateLimiter rateLimiter) {
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        String rateLimiterKey = buildRateLimiterKey(key, joinPoint);
        List<String> keys = Collections.singletonList(rateLimiterKey);
        try {
            Long number = redisTemplate.execute(limitScript, keys, count, time);
            // 如果请求次数超过限流次数抛出异常提示
            if (number==null || number.intValue() > count) {
                throw new BizException(CommonResponseCodeEnum.VERIFICATION_CODE_SEND_FREQUENTLY);
            }
            log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key);
        } catch (BizException e) {
            throw e;
        } catch (Exception e) {
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }
    public String buildRateLimiterKey(String key, JoinPoint point) {
        StringBuilder stringBuilder = new StringBuilder(key);
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();
        String requestURI = request.getRequestURI().replace("/","-").substring(1);
        String ipAddress = IpUtil.getIpAddress(request).replace(":","-");
        stringBuilder.append(requestURI).append(":").append(ipAddress);
        return stringBuilder.toString();
    }

}
